# Making Figure
rm(list=ls())
library(latex2exp) # ver 0.4.0
load(file = "../data/Adj_sim.rdata")
source("Twitter/simulation_setup.R")
load("../results/twitter_sim_output.rdata")

# 1. Linear Case
est_G_l <- output$est_G_l_f[[1]]
coef_mat_l <- output$coef_mat_l_f[[1]]

# TRUE ANSE
p1 <-  0.6; p2 <- 0.4
true_ANSE <- rep((p1 - p2)*true_G, length(overlap_use))

est_G_mean <- unlist(lapply(est_G_l, function(x)  mean((p1-p2)*x) ))
est_G_se   <- unlist(lapply(est_G_l, function(x) (p1-p2)*sd(x)))

# Parametric sensitivity analysis 
coef_U_b <- coef_U_l
para_cor <- overlap_use*coef_U_b*(p1-p2)
est_G_para_mean  <-  est_G_mean - para_cor
est_G_para_se <- est_G_se

# Nonparametric sensitivity analysis
p_own <- 0.5; pu1 <- 0.8;  pu2 <- 0.2
# approximating the MR_UY
MR_UY <- (intercept_l + true_G*p2 + coef_U_l*pu1)/(intercept_l + true_G*p2 +  coef_U_l*pu2)
RR_GU <- (overlap_use*p1  + (1-overlap_use)*p_own)/(overlap_use*p2  + (1-overlap_use)*p_own)
B <- (MR_UY*RR_GU)/(MR_UY + RR_GU - 1)

non_lower <- non_upper <- c()
non_lower_se <- non_upper_se <- c()
for(z in 1:length(overlap_use)){
  
  exp_G <- coef_mat_l[[z]][, 3]
  m_high <- coef_mat_l[[z]][, 1] + exp_G*p1
  m_low  <- coef_mat_l[[z]][, 1] + exp_G*p2
  
  non_lower[z] <- mean(m_high/B[z] - B[z]*m_low)
  non_lower_se[z] <- sd(m_high/B[z] - B[z]*m_low)
  non_upper[z] <- mean(m_high*B[z] - m_low/B[z])
  non_upper_se[z] <- sd(m_high*B[z] - m_low/B[z])
}

# 2.Interactive Model 
est_G_l2 <- output$est_G_l_f[[2]]
coef_mat_l2 <- output$coef_mat_l_f[[2]]

# TRUE ANSE
p1 <-  0.6; p2 <- 0.4
true_ANSE2 <- (p1 - p2)*true_G +  coef_U2*(p1-p2)*(overlap_use*p2  + (1-overlap_use)*0.5)

est_G_mean2  <- unlist(lapply(est_G_l2, function(x)  mean((p1-p2)*x) ))
est_G_se2   <- unlist(lapply(est_G_l2, function(x) (p1-p2)*sd(x)))

# Parametric sensitivity analysis 
coef_U_b2 <- coef_U1
para_cor2 <- overlap_use*coef_U_b2*(p1-p2)
est_G_para_mean2  <-  est_G_mean2 - para_cor2
est_G_para_se2 <- est_G_se2

# Nonparametric sensitivity analysis 
# approximating the MR_UY
MR_UY2 <- (intercept + true_G*p2 + coef_U2*p2*pu1 + coef_U1*pu1)/(intercept + true_G*p2 + coef_U2*p2*pu2 + coef_U1*pu2)
RR_GU2 <- (overlap_use*p1  + (1-overlap_use)*p_own)/(overlap_use*p2  + (1-overlap_use)*p_own)
B2 <- (MR_UY2*RR_GU2)/(MR_UY2 + RR_GU2 - 1)

non_lower2 <- non_upper2 <- c()
non_lower_se2 <- non_upper_se2 <- c()
for(z in 1:length(overlap_use)){
  m_high2 <- coef_mat_l2[[z]][, 1] + coef_mat_l2[[z]][, 3]*p1
  m_low2 <- coef_mat_l2[[z]][, 1] + coef_mat_l2[[z]][, 3]*p2
  
  non_lower2[z] <- mean(m_high2/B2[z] - B2[z]*m_low2)
  non_lower_se2[z] <- sd(m_high2/B2[z] - B2[z]*m_low2)
  non_upper2[z] <- mean(m_high2*B2[z] - m_low2/B2[z])
  non_upper_se2[z] <- sd(m_high2*B2[z] - m_low2/B2[z])
}

# ######################
# Plot 
# ######################
# Raw estimates 
ylim_use  <- ylim_use2 <- c(-0.1, 0.85)
ylim_use22  <- ylim_use02 <- c(0.1,  1.1)
xlim_use <- c(-0.1, 0.7)

pdf("../results/Fig_2.pdf",  height = 8, width = 12.5)
par(mfrow = c(2,3), mar =  c(4, 2, 2, 2), oma = c(0, 5, 1, 0))
# 1. Linear Model
plot(overlap_use, est_G_mean, ylim =  ylim_use, 
     pch = NA, xlim = xlim_use, 
     xlab = TeX("$\\pi_{GU}$: Overlap"), ylab = "",
     cex.lab = 1.75, cex.axis  =  1.75, main = "", cex.main = 1.55)
segments(overlap_use, est_G_mean - 1.96*est_G_se,
         overlap_use, est_G_mean + 1.96*est_G_se, lwd = 2)
points(overlap_use, est_G_mean, pch = 21, bg = "white", cex = 1.5, lwd = 2)
abline(h= true_ANSE[1],  col = "red", lty = 2, lwd =  2)
text(x = 0.6, y = 0.15, "True Effect", col = "red", font = 2, cex =  1.5)
mtext(side  = 2, text = "Estimated Effects", cex  = 1.15,  line =  3)
mtext(side  = 3, text = "Biased Estimator", cex  = 1.35,  font = 2, line =  1)
mtext(side  = 2, text = "Linear Additive Model", cex  = 1.25,  line =  5, font = 2)

# parametric sensitivity estimates 
plot(overlap_use, est_G_para_mean, ylim =  ylim_use, 
     pch = NA, xlim = xlim_use, 
     xlab = TeX("$\\pi_{GU}$: Overlap"), ylab = "",
     cex.lab = 1.75, cex.axis  =  1.75, main = "", cex.main = 1.55)
segments(overlap_use, est_G_para_mean - 1.96*est_G_para_se,
         overlap_use, est_G_para_mean + 1.96*est_G_para_se, lwd = 2, col = "blue")
points(overlap_use, est_G_para_mean, pch = 21, bg = "white", cex = 1.5, lwd = 2, col = "blue")
abline(h= true_ANSE[1],  col = "red", lty = 2, lwd =  2)
mtext(side  = 3, text = "Parametric Sensitivity Analysis", cex  = 1.35,  font = 2, line =  1)

# nonparametric sensitivity estimates 
plot(overlap_use, non_lower, xlim = xlim_use, ylim =  ylim_use2, pch = NA, 
     col =  "green4",  cex =  1.5, 
     xlab = TeX("$\\pi_{GU}$: Overlap"), ylab = "",
     cex.lab = 1.75, cex.axis  =  1.75, main = "", cex.main = 1.55)
points(overlap_use[1], non_lower[1], pch = 21, bg = "green4", cex = 1.5, lwd = 2, col = "green4")
arrows(overlap_use[2:4], non_lower[2:4], overlap_use[2:4], non_upper[2:4], 
       lwd = 4, code = 3, length = 0.03, angle = 90,
       col =  "green4")
arrows(overlap_use, non_lower -1.96*non_lower_se, 
       overlap_use, non_upper + 1.96*non_upper_se, 
       lwd = 2, code = 3, length = 0.03, angle = 90,
       col =  "green4")
abline(h= true_ANSE[1],  col = "red", lty = 2, lwd =  2)
mtext(side  = 3, text = "Nonparametric Sensitivity Analysis", cex  = 1.35,  font = 2, line =  1)

# 2. Interactive Model
plot(overlap_use, est_G_mean2, ylim =  ylim_use02, 
     pch = NA, xlim = xlim_use, 
     xlab = TeX("$\\pi_{GU}$: Overlap"), ylab = "",
     cex.lab = 1.75, cex.axis  =  1.75, main = "", cex.main = 1.55)
segments(overlap_use, est_G_mean2 - 1.96*est_G_se2,
         overlap_use, est_G_mean2 + 1.96*est_G_se2, lwd = 2)
points(overlap_use, est_G_mean2, pch = 21, bg = "white", cex = 1.5, lwd = 2)
abline(h= true_ANSE2[1],  col = "red", lty = 2, lwd =  2)
text(x = 0.6, y = 0.33, "True Effect", col = "red", font = 2, cex =  1.5)
mtext(side  = 2, text = "Estimated Effects", cex  = 1.15,  line =  3)
mtext(side  = 2, text = "Interactive Model", cex  = 1.25,  line =  5, font = 2)

# parametric sensitivity estimates 
plot(overlap_use, est_G_para_mean2, ylim =  ylim_use02, 
     pch = NA, xlim = xlim_use, 
     xlab = TeX("$\\pi_{GU}$: Overlap"), ylab = "",
     cex.lab = 1.75, cex.axis  =  1.75, main = "", cex.main = 1.55)
segments(overlap_use, est_G_para_mean2 - 1.96*est_G_para_se2,
         overlap_use, est_G_para_mean2 + 1.96*est_G_para_se2, lwd = 2, col = "blue")
points(overlap_use, est_G_para_mean2, pch = 21, bg = "white", cex = 1.5, lwd = 2, col = "blue")
abline(h= true_ANSE2[1],  col = "red", lty = 2, lwd =  2)

# nonparametric sensitivity estimates 
plot(overlap_use, non_lower2, xlim = xlim_use, ylim =  ylim_use22, pch = NA, 
     col =  "green4",  cex =  1.5, 
     xlab = TeX("$\\pi_{GU}$: Overlap"), ylab = "",
     cex.lab = 1.75, cex.axis  =  1.75, main = "", cex.main = 1.55)
points(overlap_use[1], non_lower2[1], pch = 21, bg = "green4", cex = 1.5, lwd = 2, col = "green4")
arrows(overlap_use[2:4], non_lower2[2:4], overlap_use[2:4], non_upper2[2:4], 
       lwd = 4, code = 3, length = 0.03, angle = 90,
       col =  "green4")
arrows(overlap_use, non_lower2 -1.96*non_lower_se2, 
       overlap_use, non_upper2 + 1.96*non_upper_se2, 
       lwd = 2, code = 3, length = 0.03, angle = 90,
       col =  "green4")
abline(h= true_ANSE2[1],  col = "red", lty = 2, lwd =  2)
dev.off()